use std::rc::Rc;

use crate::{ math::Interval, scene::Hitables, HitRecord, Hitable, Ray, AABB };

/// BVH node
/// 
pub struct BVHNode {
    // left child of BVH node
    l: Rc<Box<dyn Hitable>>,
    // right child of BVH node
    r: Rc<Box<dyn Hitable>>,
    // aabb of BVH node
    aabb: AABB,
}

impl BVHNode {
    pub fn new(mut hitables: Vec<Box<dyn Hitable>>) -> Self{
        assert!(!hitables.is_empty());

        let mut aabb = AABB::default();

        for hitable in &hitables {
            aabb = AABB::from_two_aabb(
                &aabb,
                hitable.bounding_box(),
            );
        }

        // opt, use longest axis to build bvh tree
        let axis = aabb.longest_axis();

        let (l, r) = match hitables.len() {
            // just one element in hitables vector
            1 => {
                let v0 = Rc::new(hitables.remove(0));

                (v0.clone(), v0)
            },
            // just two element in hitables vector
            2 => {
                let v1 = Rc::new(hitables.remove(1));
                let v0 = Rc::new(hitables.remove(0));

                (v0, v1)
            },
            _ => {
                hitables.sort_by(|a, b| {
                    let aabb_a = a.bounding_box();
                    let aabb_b = b.bounding_box();

                    let a = aabb_a.axis_interval(axis).min;
                    let b = aabb_b.axis_interval(axis).min;

                    a.partial_cmp(&b).unwrap()
                });

                let mid = hitables.len() / 2;

                let b = hitables.drain(mid..).collect();
                let a = hitables;

                let l = Rc::new(Box::new(BVHNode::new(a)) as Box<dyn Hitable>);
                let r = Rc::new(Box::new(BVHNode::new(b)) as Box<dyn Hitable>);

                (l, r)
            }
        };

        let aabb = AABB::from_two_aabb(
            l.bounding_box(),
            r.bounding_box(),
        );

        Self { l, r, aabb }
    }
}

impl From<Hitables> for BVHNode {
    fn from(hitables: Hitables) -> Self {
        Self::new(hitables.objects)
    }
}

impl Hitable for BVHNode {
    fn hit(&self, ray: &Ray, interval: Interval) -> Option<HitRecord> {
        if !self.aabb.hit(ray, interval) {
            return None;
        }

        let hit_l = self.l.hit(ray, interval);

        let t_max = hit_l
            .as_ref()
            .map_or(interval.max, |hit| hit.t);

        let hit_r = self.r.hit(ray, Interval::new(
            interval.min,
            t_max,
        ));

        hit_r.or(hit_l)
    }

    fn bounding_box(&self) -> &AABB {
        &self.aabb
    }
}